{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Copyright 2025 Apache Software Foundation\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ],
   "id": "f19821bfb84682cd"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Burr Demo - observing state\n",
    "# & traveling back in time!\n",
    "\n",
    "\n",
    "<img src=\"https://github.com/user-attachments/assets/2ab9b499-7ca2-4ae9-af72-ccc775f30b4e\" width=\"100\" align=\"left\" /> +\n",
    "<img src=\"https://cdn.mos.cms.futurecdn.net/VgGxJABA8DcfAMpPPwdv6a.jpg\" width=\"200\" align=\"center\"/>\n",
    "\n",
    "[https://github.com/dagworks-inc/burr](https://github.com/dagworks-inc/burr) by DAGWorks Inc. (YCW23 & StartX).\n",
    "\n",
    "Take🏠:\n",
    "\n",
    " - high level what is Burr\n",
    " - what you can do with Burr (observing state & being able to debug a particular point in time)\n",
    " - watch a walkthrough of this [notebook here](https://youtu.be/hqutVJyd3TI)."
   ],
   "id": "a18d66ec7ff9b666"
  },
  {
   "cell_type": "markdown",
   "id": "81128a911f2cbeee",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Agentic Problems"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff9f53c-f080-4d05-8b31-70db743f973d",
   "metadata": {},
   "source": [
    "## 1. Why did this LLM call fail?\n",
    "## 2. Oh crap my code broke, why?\n",
    "## 3. Things went off the rails, but where?\n",
    "## 4. etc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d725b4dd413d01e8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Monitoring FTW, right?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92cea7b3-1d27-4820-90b5-7aba9026fdfe",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Well but ... monitoring doesn't help you debug & complete your dev loop\n",
    "\n",
    "## 1. How do I debug that quickly?\n",
    "\n",
    "## 2. How do I fix the inputs/code, and restart my agent?\n",
    "\n",
    "## 3. What if my agent was 20+ steps in … do I have to restart from step 0? or can I go to a specific point in time?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f602812d90d20d69",
   "metadata": {},
   "source": [
    "# Solution: Burr\n",
    "(Complements our other framework [Hamilton](https://github.com/dagWorks-Inc/hamilton))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3422aed8-323f-45dd-ade1-b1f0906c35f3",
   "metadata": {},
   "source": [
    "## 1. Agent application is modeled as State + Actions --> Graph\n",
    "Straightforward multi-modal example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7953fe759ae0f510",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-30T21:57:38.408541Z",
     "start_time": "2024-09-30T21:57:36.086836Z"
    }
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from IPython.core.display import HTML \n",
    "import openai\n",
    "\n",
    "from burr.core import ApplicationBuilder, State, default, graph, when\n",
    "from burr.core.action import action\n",
    "from burr.tracking import LocalTrackingClient\n",
    "\n",
    "MODES = {\n",
    "    \"answer_question\": \"text\",\n",
    "    \"generate_image\": \"image\",\n",
    "    \"generate_code\": \"code\",\n",
    "    \"unknown\": \"text\",\n",
    "}\n",
    "\n",
    "\n",
    "@action(reads=[], writes=[\"chat_history\", \"prompt\"])\n",
    "def process_prompt(state: State, prompt: str) -> State:\n",
    "    result = {\"chat_item\": {\"role\": \"user\", \"content\": prompt, \"type\": \"text\"}}\n",
    "    state = state.append(chat_history=result[\"chat_item\"])\n",
    "    state = state.update(prompt=prompt)\n",
    "    return state\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\"], writes=[\"mode\"])\n",
    "def choose_mode(state: State) -> State:\n",
    "    prompt = (\n",
    "        f\"You are a chatbot. You've been prompted this: {state['prompt']}. \"\n",
    "        f\"You have the capability of responding in the following modes: {', '.join(MODES)}. \"\n",
    "        \"Please respond with *only* a single word representing the mode that most accurately \"\n",
    "        \"corresponds to the prompt. Fr instance, if the prompt is 'draw a picture of a cat', \"\n",
    "        \"the mode would be 'generate_image'. If the prompt is \"\n",
    "        \"'what is the capital of France', the mode would be 'answer_question'.\"\n",
    "        \"If none of these modes apply, please respond with 'unknown'.\"\n",
    "    )\n",
    "\n",
    "    llm_result = openai.Client().chat.completions.create(\n",
    "        model=\"gpt-4\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "    )\n",
    "    content = llm_result.choices[0].message.content\n",
    "    mode = content.lower()\n",
    "    if mode not in MODES:\n",
    "        mode = \"unknown\"\n",
    "    result = {\"mode\": mode}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\"], writes=[\"response\"])\n",
    "def prompt_for_more(state: State) -> State:\n",
    "    result = {\n",
    "        \"response\": {\n",
    "            \"content\": \"None of the response modes I support apply to your question. \"\n",
    "                       \"Please clarify?\",\n",
    "            \"type\": \"text\",\n",
    "            \"role\": \"assistant\",\n",
    "        }\n",
    "    }\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def chat_response(\n",
    "        state: State, prepend_prompt: str, model: str = \"gpt-3.5-turbo\"\n",
    ") -> State:\n",
    "    \n",
    "    chat_history = copy.deepcopy(state[\"chat_history\"])\n",
    "    chat_history[-1][\"content\"] = f\"{prepend_prompt}: {chat_history[-1]['content']}\"\n",
    "    chat_history_api_format = [\n",
    "        {\n",
    "            \"role\": chat[\"role\"],\n",
    "            \"content\": chat[\"content\"],\n",
    "        }\n",
    "        for chat in chat_history\n",
    "    ]\n",
    "    client = openai.Client()\n",
    "    result = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=chat_history_api_format,\n",
    "    )\n",
    "    text_response = result.choices[0].message.content\n",
    "    result = {\"response\": {\"content\": text_response, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def image_response(state: State, model: str = \"dall-e-2\") -> State:\n",
    "    \"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\n",
    "    # raise ValueError(\"Demo error\")\n",
    "    client = openai.Client()\n",
    "    result = client.images.generate(\n",
    "        model=model, prompt=state[\"prompt\"], size=\"1024x1024\", quality=\"standard\", n=1\n",
    "    )\n",
    "    image_url = result.data[0].url\n",
    "    result = {\"response\": {\"content\": image_url, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"response\", \"mode\"], writes=[\"chat_history\"])\n",
    "def response(state: State) -> State:\n",
    "    # you'd do something specific here based on prior state\n",
    "    result = {\"chat_item\": state[\"response\"]}\n",
    "    return state.append(chat_history=result[\"chat_item\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "add99d139d3d72a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:23:11.529415Z",
     "start_time": "2024-07-16T18:23:11.008552Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"563pt\" height=\"307pt\"\n",
       " viewBox=\"0.00 0.00 563.47 307.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 303)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-303 559.47,-303 559.47,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M467.1,-233.4C467.1,-233.4 426.75,-233.4 426.75,-233.4 420.75,-233.4 414.75,-227.4 414.75,-221.4 414.75,-221.4 414.75,-208.8 414.75,-208.8 414.75,-202.8 420.75,-196.8 426.75,-196.8 426.75,-196.8 467.1,-196.8 467.1,-196.8 473.1,-196.8 479.1,-202.8 479.1,-208.8 479.1,-208.8 479.1,-221.4 479.1,-221.4 479.1,-227.4 473.1,-233.4 467.1,-233.4\"/>\n",
       "<text text-anchor=\"middle\" x=\"446.93\" y=\"-209.3\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M369.6,-167.8C369.6,-167.8 290.25,-167.8 290.25,-167.8 284.25,-167.8 278.25,-161.8 278.25,-155.8 278.25,-155.8 278.25,-143.2 278.25,-143.2 278.25,-137.2 284.25,-131.2 290.25,-131.2 290.25,-131.2 369.6,-131.2 369.6,-131.2 375.6,-131.2 381.6,-137.2 381.6,-143.2 381.6,-143.2 381.6,-155.8 381.6,-155.8 381.6,-161.8 375.6,-167.8 369.6,-167.8\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-143.7\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M414.33,-196.38C401.28,-189.29 386.06,-181.01 372.17,-173.46\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"374.13,-170.54 363.67,-168.84 370.79,-176.69 374.13,-170.54\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"497.48,-299 396.38,-299 396.38,-262.4 497.48,-262.4 497.48,-299\"/>\n",
       "<text text-anchor=\"middle\" x=\"446.93\" y=\"-274.9\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M446.93,-261.98C446.93,-256.75 446.93,-250.87 446.93,-245.12\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"450.43,-245.23 446.93,-235.23 443.43,-245.23 450.43,-245.23\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M107.85,-102.2C107.85,-102.2 12,-102.2 12,-102.2 6,-102.2 0,-96.2 0,-90.2 0,-90.2 0,-77.6 0,-77.6 0,-71.6 6,-65.6 12,-65.6 12,-65.6 107.85,-65.6 107.85,-65.6 113.85,-65.6 119.85,-71.6 119.85,-77.6 119.85,-77.6 119.85,-90.2 119.85,-90.2 119.85,-96.2 113.85,-102.2 107.85,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"59.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M277.8,-136.78C239.73,-128.16 186.56,-115.98 131.21,-102.73\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"132.13,-99.35 121.59,-100.42 130.5,-106.16 132.13,-99.35\"/>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M237.73,-102.2C237.73,-102.2 150.12,-102.2 150.12,-102.2 144.12,-102.2 138.12,-96.2 138.12,-90.2 138.12,-90.2 138.12,-77.6 138.12,-77.6 138.12,-71.6 144.12,-65.6 150.12,-65.6 150.12,-65.6 237.73,-65.6 237.73,-65.6 243.73,-65.6 249.73,-71.6 249.73,-77.6 249.73,-77.6 249.73,-90.2 249.73,-90.2 249.73,-96.2 243.73,-102.2 237.73,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M292.04,-130.78C276.42,-123.48 258.14,-114.93 241.62,-107.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"243.52,-104.23 232.98,-103.16 240.55,-110.57 243.52,-104.23\"/>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M379.73,-102.2C379.73,-102.2 280.12,-102.2 280.12,-102.2 274.12,-102.2 268.12,-96.2 268.12,-90.2 268.12,-90.2 268.12,-77.6 268.12,-77.6 268.12,-71.6 274.12,-65.6 280.12,-65.6 280.12,-65.6 379.73,-65.6 379.73,-65.6 385.73,-65.6 391.73,-71.6 391.73,-77.6 391.73,-77.6 391.73,-90.2 391.73,-90.2 391.73,-96.2 385.73,-102.2 379.73,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M329.93,-130.78C329.93,-125.55 329.93,-119.67 329.93,-113.92\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"333.43,-114.03 329.93,-104.03 326.43,-114.03 333.43,-114.03\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M524.23,-102.2C524.23,-102.2 421.62,-102.2 421.62,-102.2 415.62,-102.2 409.62,-96.2 409.62,-90.2 409.62,-90.2 409.62,-77.6 409.62,-77.6 409.62,-71.6 415.62,-65.6 421.62,-65.6 421.62,-65.6 524.23,-65.6 524.23,-65.6 530.23,-65.6 536.23,-71.6 536.23,-77.6 536.23,-77.6 536.23,-90.2 536.23,-90.2 536.23,-96.2 530.23,-102.2 524.23,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"472.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M369.76,-130.78C386.34,-123.41 405.77,-114.77 423.26,-106.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"424.21,-110.4 431.92,-103.14 421.36,-104 424.21,-110.4\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M356.48,-36.6C356.48,-36.6 303.38,-36.6 303.38,-36.6 297.38,-36.6 291.38,-30.6 291.38,-24.6 291.38,-24.6 291.38,-12 291.38,-12 291.38,-6 297.38,0 303.38,0 303.38,0 356.48,0 356.48,0 362.48,0 368.48,-6 368.48,-12 368.48,-12 368.48,-24.6 368.48,-24.6 368.48,-30.6 362.48,-36.6 356.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M120.12,-67.74C123.1,-67.01 126.04,-66.29 128.93,-65.6 180.32,-53.24 239.28,-39.74 279.99,-30.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"280.51,-33.99 289.49,-28.37 278.96,-27.17 280.51,-33.99\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"241.85,-167.8 146,-167.8 146,-131.2 241.85,-131.2 241.85,-167.8\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.93\" y=\"-143.7\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M156.6,-130.78C141.35,-123.55 123.54,-115.09 107.38,-107.42\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"108.95,-104.3 98.42,-103.17 105.95,-110.62 108.95,-104.3\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M193.93,-130.78C193.93,-125.55 193.93,-119.67 193.93,-113.92\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"197.43,-114.03 193.93,-104.03 190.43,-114.03 197.43,-114.03\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M231.81,-130.78C247.43,-123.48 265.71,-114.93 282.23,-107.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"283.3,-110.57 290.87,-103.16 280.33,-104.23 283.3,-110.57\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M231.81,-65.18C247.43,-57.88 265.71,-49.33 282.23,-41.61\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"283.3,-44.97 290.87,-37.56 280.33,-38.63 283.3,-44.97\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M329.93,-65.18C329.93,-59.95 329.93,-54.07 329.93,-48.32\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"333.43,-48.43 329.93,-38.43 326.43,-48.43 333.43,-48.43\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M433.09,-65.18C416.24,-57.69 396.44,-48.88 378.71,-41\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"380.5,-37.96 369.94,-37.1 377.66,-44.36 380.5,-37.96\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M368.88,-21.97C424.73,-26.7 523.21,-38.72 544.92,-65.6 578.23,-106.85 522.83,-159.53 482.7,-189.78\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"480.85,-186.79 474.86,-195.53 485,-192.44 480.85,-186.79\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x12f4f5510>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Built the graph.\n",
    "base_graph = (\n",
    "    graph.GraphBuilder()\n",
    "    .with_actions(\n",
    "        # these are the \"nodes\" \n",
    "        prompt=process_prompt,\n",
    "        decide_mode=choose_mode,\n",
    "        generate_image=image_response,\n",
    "        generate_code=chat_response.bind(\n",
    "            prepend_prompt=\"Please respond with *only* code and no other text (at all) to the following:\",\n",
    "        ),\n",
    "        answer_question=chat_response.bind(\n",
    "            prepend_prompt=\"Please answer the following question:\",\n",
    "        ),\n",
    "        prompt_for_more=prompt_for_more,\n",
    "        response=response,\n",
    "    )\n",
    "    .with_transitions(\n",
    "        # these are the edges between nodes, based on state.\n",
    "        (\"prompt\", \"decide_mode\", default),\n",
    "        (\"decide_mode\", \"generate_image\", when(mode=\"generate_image\")),\n",
    "        (\"decide_mode\", \"generate_code\", when(mode=\"generate_code\")),\n",
    "        (\"decide_mode\", \"answer_question\", when(mode=\"answer_question\")),\n",
    "        (\"decide_mode\", \"prompt_for_more\", default),\n",
    "        (\n",
    "            [\"generate_image\", \"answer_question\", \"generate_code\", \"prompt_for_more\"],\n",
    "            \"response\",\n",
    "        ),\n",
    "        (\"response\", \"prompt\", default),\n",
    "    )\n",
    "    .build()\n",
    ")\n",
    "base_graph.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b26ad97a43444",
   "metadata": {},
   "source": [
    "## 2. Build application --> built in checkpointing & tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4f0bf205944272ff",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:27:12.136291Z",
     "start_time": "2024-07-16T18:27:11.518066Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"576pt\" height=\"331pt\"\n",
       " viewBox=\"0.00 0.00 575.58 330.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 326.5)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-326.5 571.58,-326.5 571.58,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M509.1,-254.9C509.1,-254.9 468.75,-254.9 468.75,-254.9 462.75,-254.9 456.75,-248.9 456.75,-242.9 456.75,-242.9 456.75,-230.3 456.75,-230.3 456.75,-224.3 462.75,-218.3 468.75,-218.3 468.75,-218.3 509.1,-218.3 509.1,-218.3 515.1,-218.3 521.1,-224.3 521.1,-230.3 521.1,-230.3 521.1,-242.9 521.1,-242.9 521.1,-248.9 515.1,-254.9 509.1,-254.9\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-230.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M433.6,-187.3C433.6,-187.3 354.25,-187.3 354.25,-187.3 348.25,-187.3 342.25,-181.3 342.25,-175.3 342.25,-175.3 342.25,-162.7 342.25,-162.7 342.25,-156.7 348.25,-150.7 354.25,-150.7 354.25,-150.7 433.6,-150.7 433.6,-150.7 439.6,-150.7 445.6,-156.7 445.6,-162.7 445.6,-162.7 445.6,-175.3 445.6,-175.3 445.6,-181.3 439.6,-187.3 433.6,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"393.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M463.47,-218.02C452.82,-210.67 440.27,-202 428.79,-194.07\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"430.92,-191.29 420.7,-188.49 426.94,-197.05 430.92,-191.29\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"539.48,-322.5 438.38,-322.5 438.38,-285.9 539.48,-285.9 539.48,-322.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-298.4\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M488.93,-285.62C488.93,-279.74 488.93,-273.02 488.93,-266.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"492.43,-266.74 488.93,-256.74 485.43,-266.74 492.43,-266.74\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M117.85,-104.2C117.85,-104.2 22,-104.2 22,-104.2 16,-104.2 10,-98.2 10,-92.2 10,-92.2 10,-79.6 10,-79.6 10,-73.6 16,-67.6 22,-67.6 22,-67.6 117.85,-67.6 117.85,-67.6 123.85,-67.6 129.85,-73.6 129.85,-79.6 129.85,-79.6 129.85,-92.2 129.85,-92.2 129.85,-98.2 123.85,-104.2 117.85,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"69.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M341.99,-167.51C269.03,-165.96 140.74,-159.76 101.68,-135.7 93.45,-130.63 86.93,-122.54 82.01,-114.39\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"85.11,-112.77 77.31,-105.6 78.93,-116.07 85.11,-112.77\"/>\n",
       "<text text-anchor=\"middle\" x=\"164.55\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M247.73,-104.2C247.73,-104.2 160.12,-104.2 160.12,-104.2 154.12,-104.2 148.12,-98.2 148.12,-92.2 148.12,-92.2 148.12,-79.6 148.12,-79.6 148.12,-73.6 154.12,-67.6 160.12,-67.6 160.12,-67.6 247.73,-67.6 247.73,-67.6 253.73,-67.6 259.73,-73.6 259.73,-79.6 259.73,-79.6 259.73,-92.2 259.73,-92.2 259.73,-98.2 253.73,-104.2 247.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"203.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M342.02,-163.6C313.66,-159.45 278.75,-151.38 250.93,-135.7 240.73,-129.96 231.37,-121.31 223.72,-112.88\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"226.55,-110.81 217.4,-105.47 221.22,-115.35 226.55,-110.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"309.43\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M389.73,-104.2C389.73,-104.2 290.12,-104.2 290.12,-104.2 284.12,-104.2 278.12,-98.2 278.12,-92.2 278.12,-92.2 278.12,-79.6 278.12,-79.6 278.12,-73.6 284.12,-67.6 290.12,-67.6 290.12,-67.6 389.73,-67.6 389.73,-67.6 395.73,-67.6 401.73,-73.6 401.73,-79.6 401.73,-79.6 401.73,-92.2 401.73,-92.2 401.73,-98.2 395.73,-104.2 389.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M387.8,-150.49C383.99,-140.82 378.56,-128.85 371.93,-119.2 370.51,-117.14 368.95,-115.09 367.31,-113.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"370.06,-110.91 360.78,-105.8 364.85,-115.58 370.06,-110.91\"/>\n",
       "<text text-anchor=\"middle\" x=\"444.8\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M534.23,-104.2C534.23,-104.2 431.62,-104.2 431.62,-104.2 425.62,-104.2 419.62,-98.2 419.62,-92.2 419.62,-92.2 419.62,-79.6 419.62,-79.6 419.62,-73.6 425.62,-67.6 431.62,-67.6 431.62,-67.6 534.23,-67.6 534.23,-67.6 540.23,-67.6 546.23,-73.6 546.23,-79.6 546.23,-79.6 546.23,-92.2 546.23,-92.2 546.23,-98.2 540.23,-104.2 534.23,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"482.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M445.99,-160.38C474.35,-155.05 505.36,-146.87 513.92,-135.7 519.37,-128.6 517.26,-120.53 512.21,-113.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"515.23,-111.19 506.16,-105.73 509.85,-115.66 515.23,-111.19\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M366.48,-36.6C366.48,-36.6 313.38,-36.6 313.38,-36.6 307.38,-36.6 301.38,-30.6 301.38,-24.6 301.38,-24.6 301.38,-12 301.38,-12 301.38,-6 307.38,0 313.38,0 313.38,0 366.48,0 366.48,0 372.48,0 378.48,-6 378.48,-12 378.48,-12 378.48,-24.6 378.48,-24.6 378.48,-30.6 372.48,-36.6 366.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M130.13,-69.77C133.1,-69.03 136.05,-68.31 138.93,-67.6 190.38,-54.97 249.33,-40.84 290.02,-31.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"290.6,-34.61 299.52,-28.89 288.98,-27.8 290.6,-34.61\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"95.85,-187.3 0,-187.3 0,-150.7 95.85,-150.7 95.85,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M38.84,-150.37C35.2,-140.89 32.75,-129.13 36.93,-119.2 37.79,-117.15 38.85,-115.16 40.06,-113.25\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"42.69,-115.57 46.15,-105.56 37.2,-111.23 42.69,-115.57\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M48.56,-150.28C49.89,-139.79 53.35,-127.06 61.93,-119.2 85.32,-97.76 101.77,-109.4 136.97,-104.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.31,-107.65 146.48,-102.34 136,-100.77 137.31,-107.65\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M96.05,-162.46C128.21,-157.75 171.03,-149.48 206.93,-135.7 220.41,-130.52 221.71,-125.03 234.93,-119.2 245.07,-114.73 256.02,-110.62 266.88,-106.94\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.83,-110.31 276.25,-103.88 265.66,-103.66 267.83,-110.31\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M240.73,-67.15C256.75,-59.42 275.73,-50.26 292.77,-42.04\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"294.24,-45.22 301.73,-37.72 291.2,-38.92 294.24,-45.22\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M339.93,-67.32C339.93,-61.44 339.93,-54.72 339.93,-48.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"343.43,-48.44 339.93,-38.44 336.43,-48.44 343.43,-48.44\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M444.23,-67.15C427.22,-59.34 407.05,-50.09 389,-41.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"390.58,-38.69 380.03,-37.7 387.66,-45.05 390.58,-38.69\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M378.9,-22.35C434.77,-27.62 533.29,-40.57 554.92,-67.6 589.49,-110.79 545,-174.31 514.04,-209.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"511.64,-207.01 507.54,-216.78 516.84,-211.7 511.64,-207.01\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<burr.core.application.Application at 0x12b7aa6e0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n",
    "# this will auto instrument the openAI client. No swapping of imports required!\n",
    "OpenAIInstrumentor().instrument()\n",
    "\n",
    "tracker = LocalTrackingClient(project=\"agent-demo\")\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker, \n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)  # tracking + checkpointing; one line 🪄.\n",
    "    .build()\n",
    ")\n",
    "app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c48e44fbf36f4d",
   "metadata": {},
   "source": [
    "## 3. Comes with a UI\n",
    "View runs in the UI; Let's run the app first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48745f4c04b9d28d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? say hi in SQL\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: SELECT 'hi';\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? draw a cat\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "********************************************************************************\n",
      "-------------------------------------------------------------------\n",
      "Oh no an error! Need help with Burr?\n",
      "Join our discord and ask for help! https://discord.gg/4FxBMyzW5n\n",
      "-------------------------------------------------------------------\n",
      "> Action: `generate_image` encountered an error!<\n",
      "> State (at time of action):\n",
      "{'__PRIOR_STEP': 'decide_mode',\n",
      " '__SEQUENCE_ID': 6,\n",
      " 'chat_history': \"[{'role': 'user', 'content': 'say hi in SQL', 'typ...\",\n",
      " 'mode': 'generate_image',\n",
      " 'prompt': 'draw a cat',\n",
      " 'response': '{\\'content\\': \"SELECT \\'hi\\';\", \\'type\\': \\'code\\', \\'role\\'...'}\n",
      "> Inputs (at time of action):\n",
      "{'prompt': 'draw a cat'}\n",
      "********************************************************************************\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/application.py\", line 812, in _step\n",
      "    result, new_state = _run_single_step_action(\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/application.py\", line 252, in _run_single_step_action\n",
      "    action.run_and_update(state, **inputs), action.name, action.schema\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/action.py\", line 639, in run_and_update\n",
      "    return self._fn(state, **self._bound_params, **run_kwargs)\n",
      "  File \"/var/folders/gv/q39lb_1s26x7gbyyypqc3dkm0000gn/T/ipykernel_31316/1354917547.py\", line 94, in image_response\n",
      "    raise ValueError(\"Demo error\")\n",
      "ValueError: Demo error\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Demo error",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mValueError\u001B[0m                                Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[4], line 5\u001B[0m\n\u001B[1;32m      3\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mquit\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;241m==\u001B[39m user_input\u001B[38;5;241m.\u001B[39mlower():\n\u001B[1;32m      4\u001B[0m     \u001B[38;5;28;01mbreak\u001B[39;00m\n\u001B[0;32m----> 5\u001B[0m last_action, action_result, app_state \u001B[38;5;241m=\u001B[39m \u001B[43mapp\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m      6\u001B[0m \u001B[43m    \u001B[49m\u001B[43mhalt_after\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m[\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mresponse\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m]\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\n\u001B[1;32m      7\u001B[0m \u001B[43m    \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m{\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mprompt\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m:\u001B[49m\u001B[43m \u001B[49m\u001B[43muser_input\u001B[49m\u001B[43m}\u001B[49m\n\u001B[1;32m      8\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m      9\u001B[0m last_message \u001B[38;5;241m=\u001B[39m app_state[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mchat_history\u001B[39m\u001B[38;5;124m\"\u001B[39m][\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m]\n\u001B[1;32m     10\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m last_message[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mtype\u001B[39m\u001B[38;5;124m'\u001B[39m] \u001B[38;5;241m==\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mimage\u001B[39m\u001B[38;5;124m'\u001B[39m:\n",
      "File \u001B[0;32m~/dagworks/burr/burr/telemetry.py:276\u001B[0m, in \u001B[0;36mcapture_function_usage.<locals>.wrapped_fn\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    273\u001B[0m \u001B[38;5;129m@functools\u001B[39m\u001B[38;5;241m.\u001B[39mwraps(call_fn)\n\u001B[1;32m    274\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mwrapped_fn\u001B[39m(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    275\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 276\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mcall_fn\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    277\u001B[0m     \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    278\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m is_telemetry_enabled():\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:616\u001B[0m, in \u001B[0;36m_call_execute_method_pre_post.__call__.<locals>.wrapper_sync\u001B[0;34m(app_self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    614\u001B[0m exc \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m    615\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 616\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mapp_self\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    617\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    618\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mcall_post(app_self, exc)\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:1168\u001B[0m, in \u001B[0;36mApplication.run\u001B[0;34m(self, halt_before, halt_after, inputs)\u001B[0m\n\u001B[1;32m   1166\u001B[0m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m:\n\u001B[1;32m   1167\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m-> 1168\u001B[0m         \u001B[38;5;28;43mnext\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mgen\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1169\u001B[0m     \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mStopIteration\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m   1170\u001B[0m         result \u001B[38;5;241m=\u001B[39m e\u001B[38;5;241m.\u001B[39mvalue\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:1111\u001B[0m, in \u001B[0;36mApplication.iterate\u001B[0;34m(self, halt_before, halt_after, inputs)\u001B[0m\n\u001B[1;32m   1108\u001B[0m prior_action: Optional[Action] \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m   1109\u001B[0m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhas_next_action():\n\u001B[1;32m   1110\u001B[0m     \u001B[38;5;66;03m# self.step will only return None if there is no next action, so we can rely on tuple unpacking\u001B[39;00m\n\u001B[0;32m-> 1111\u001B[0m     prior_action, result, state \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstep\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1112\u001B[0m     \u001B[38;5;28;01myield\u001B[39;00m prior_action, result, state\n\u001B[1;32m   1113\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_should_halt_iterate(halt_before, halt_after, prior_action):\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:616\u001B[0m, in \u001B[0;36m_call_execute_method_pre_post.__call__.<locals>.wrapper_sync\u001B[0;34m(app_self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    614\u001B[0m exc \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m    615\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 616\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mapp_self\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    617\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    618\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mcall_post(app_self, exc)\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:773\u001B[0m, in \u001B[0;36mApplication.step\u001B[0;34m(self, inputs)\u001B[0m\n\u001B[1;32m    770\u001B[0m \u001B[38;5;66;03m# we need to increment the sequence before we start computing\u001B[39;00m\n\u001B[1;32m    771\u001B[0m \u001B[38;5;66;03m# that way if we're replaying from state, we don't get stuck\u001B[39;00m\n\u001B[1;32m    772\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_increment_sequence_id()\n\u001B[0;32m--> 773\u001B[0m out \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_step\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m_run_hooks\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m)\u001B[49m\n\u001B[1;32m    774\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m out\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:826\u001B[0m, in \u001B[0;36mApplication._step\u001B[0;34m(self, inputs, _run_hooks)\u001B[0m\n\u001B[1;32m    824\u001B[0m     exc \u001B[38;5;241m=\u001B[39m e\n\u001B[1;32m    825\u001B[0m     logger\u001B[38;5;241m.\u001B[39mexception(_format_BASE_ERROR_MESSAGE(next_action, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_state, inputs))\n\u001B[0;32m--> 826\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m e\n\u001B[1;32m    827\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    828\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m _run_hooks:\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:812\u001B[0m, in \u001B[0;36mApplication._step\u001B[0;34m(self, inputs, _run_hooks)\u001B[0m\n\u001B[1;32m    810\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m    811\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m next_action\u001B[38;5;241m.\u001B[39msingle_step:\n\u001B[0;32m--> 812\u001B[0m         result, new_state \u001B[38;5;241m=\u001B[39m \u001B[43m_run_single_step_action\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    813\u001B[0m \u001B[43m            \u001B[49m\u001B[43mnext_action\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_state\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43maction_inputs\u001B[49m\n\u001B[1;32m    814\u001B[0m \u001B[43m        \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    815\u001B[0m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    816\u001B[0m         result \u001B[38;5;241m=\u001B[39m _run_function(\n\u001B[1;32m    817\u001B[0m             next_action, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_state, action_inputs, name\u001B[38;5;241m=\u001B[39mnext_action\u001B[38;5;241m.\u001B[39mname\n\u001B[1;32m    818\u001B[0m         )\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:252\u001B[0m, in \u001B[0;36m_run_single_step_action\u001B[0;34m(action, state, inputs)\u001B[0m\n\u001B[1;32m    249\u001B[0m \u001B[38;5;66;03m# TODO -- guard all reads/writes with a subset of the state\u001B[39;00m\n\u001B[1;32m    250\u001B[0m action\u001B[38;5;241m.\u001B[39mvalidate_inputs(inputs)\n\u001B[1;32m    251\u001B[0m result, new_state \u001B[38;5;241m=\u001B[39m _adjust_single_step_output(\n\u001B[0;32m--> 252\u001B[0m     \u001B[43maction\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun_and_update\u001B[49m\u001B[43m(\u001B[49m\u001B[43mstate\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m)\u001B[49m, action\u001B[38;5;241m.\u001B[39mname, action\u001B[38;5;241m.\u001B[39mschema\n\u001B[1;32m    253\u001B[0m )\n\u001B[1;32m    254\u001B[0m _validate_result(result, action\u001B[38;5;241m.\u001B[39mname, action\u001B[38;5;241m.\u001B[39mschema)\n\u001B[1;32m    255\u001B[0m out \u001B[38;5;241m=\u001B[39m result, _state_update(state, new_state)\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/action.py:639\u001B[0m, in \u001B[0;36mFunctionBasedAction.run_and_update\u001B[0;34m(self, state, **run_kwargs)\u001B[0m\n\u001B[1;32m    638\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mrun_and_update\u001B[39m(\u001B[38;5;28mself\u001B[39m, state: State, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mrun_kwargs) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m \u001B[38;5;28mtuple\u001B[39m[\u001B[38;5;28mdict\u001B[39m, State]:\n\u001B[0;32m--> 639\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_fn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mstate\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_bound_params\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mrun_kwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "Cell \u001B[0;32mIn[1], line 94\u001B[0m, in \u001B[0;36mimage_response\u001B[0;34m(state, model)\u001B[0m\n\u001B[1;32m     91\u001B[0m \u001B[38;5;129m@action\u001B[39m(reads\u001B[38;5;241m=\u001B[39m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mprompt\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mchat_history\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mmode\u001B[39m\u001B[38;5;124m\"\u001B[39m], writes\u001B[38;5;241m=\u001B[39m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mresponse\u001B[39m\u001B[38;5;124m\"\u001B[39m])\n\u001B[1;32m     92\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mimage_response\u001B[39m(state: State, model: \u001B[38;5;28mstr\u001B[39m \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mdall-e-2\u001B[39m\u001B[38;5;124m\"\u001B[39m) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m State:\n\u001B[1;32m     93\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\u001B[39;00m\n\u001B[0;32m---> 94\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mDemo error\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[1;32m     95\u001B[0m     client \u001B[38;5;241m=\u001B[39m openai\u001B[38;5;241m.\u001B[39mClient()\n\u001B[1;32m     96\u001B[0m     result \u001B[38;5;241m=\u001B[39m client\u001B[38;5;241m.\u001B[39mimages\u001B[38;5;241m.\u001B[39mgenerate(\n\u001B[1;32m     97\u001B[0m         model\u001B[38;5;241m=\u001B[39mmodel, prompt\u001B[38;5;241m=\u001B[39mstate[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mprompt\u001B[39m\u001B[38;5;124m\"\u001B[39m], size\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m1024x1024\u001B[39m\u001B[38;5;124m\"\u001B[39m, quality\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mstandard\u001B[39m\u001B[38;5;124m\"\u001B[39m, n\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m1\u001B[39m\n\u001B[1;32m     98\u001B[0m     )\n",
      "\u001B[0;31mValueError\u001B[0m: Demo error"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efe2430b34bf8e5",
   "metadata": {},
   "source": [
    "[Open the UI - http://localhost:7241/](http://localhost:7241)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bfbd6bf99f7b14",
   "metadata": {},
   "source": [
    "## But something broke / I want to debug\n",
    "\n",
    "Use:\n",
    "\n",
    "* Application ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f5588457208604d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "app_id = \"3ad172ca-b9dc-4419-8948-922909d72bd4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b4abab5aa575628d",
   "metadata": {},
   "outputs": [],
   "source": [
    "resumed_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)\n",
    "    .with_identifiers(app_id=app_id)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4b21ade9-051c-453d-b3fd-f191f970c7b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'say hi in SQL', 'type': 'text'},\n",
       " {'content': \"SELECT 'hi';\", 'type': 'code', 'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'draw a cat', 'type': 'text'}]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resumed_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2186ef09eaff58d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://oaidalleapiprodscus.blob.core.windows.net/private/org-UENfV68Ccazb3rdeYDLRWmoF/user-KU9L39oZGkNLkg6ZfSLMc4AV/img-8yV9ub451PVi4nGXeQdoRH9D.png?st=2024-10-02T04%3A13%3A59Z&se=2024-10-02T06%3A13%3A59Z&sp=r&sv=2024-08-04&sr=b&rscd=inline&rsct=image/png&skoid=d505667d-d6c1-4a0a-bac7-5c84a87759f8&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skt=2024-10-02T01%3A26%3A51Z&ske=2024-10-03T01%3A26%3A51Z&sks=b&skv=2024-08-04&sig=QR5BxlPBqbYVbpvFsPDwLKyJ2uj/6lbR06QbFry3bbU%3D\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = resumed_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bec09d1e740a30d",
   "metadata": {},
   "source": [
    "## Actually what if I want to go back to a certain point in time?\n",
    "\n",
    "* Fork: Start with state from any checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c008089be309c098",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:49:55.910255Z",
     "start_time": "2024-07-16T18:49:55.906299Z"
    }
   },
   "outputs": [],
   "source": [
    "app_id = \"3ad172ca-b9dc-4419-8948-922909d72bd4\"\n",
    "sequence_id = 4\n",
    "# partition_key = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a8c5631cfebe8e48",
   "metadata": {},
   "outputs": [],
   "source": [
    "forked_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph) # this could be different...\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "        fork_from_app_id=app_id,\n",
    "        fork_from_sequence_id=sequence_id,\n",
    "        # fork_from_partition_key=partition_key\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2a725130e0532f0c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:58:04.223560Z",
     "start_time": "2024-07-16T18:58:04.218274Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'say hi in SQL', 'type': 'text'},\n",
       " {'content': \"SELECT 'hi';\", 'type': 'code', 'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'draw a cat', 'type': 'text'}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show prior forked state\n",
    "forked_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e27a17968c62aa98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://oaidalleapiprodscus.blob.core.windows.net/private/org-UENfV68Ccazb3rdeYDLRWmoF/user-KU9L39oZGkNLkg6ZfSLMc4AV/img-9x8yPX8sQTne7qUpLpU788i3.png?st=2024-10-02T04%3A15%3A54Z&se=2024-10-02T06%3A15%3A54Z&sp=r&sv=2024-08-04&sr=b&rscd=inline&rsct=image/png&skoid=d505667d-d6c1-4a0a-bac7-5c84a87759f8&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skt=2024-10-02T01%3A44%3A34Z&ske=2024-10-03T01%3A44%3A34Z&sks=b&skv=2024-08-04&sig=U%2B0tHC1QRz%2BZx0qbL/v3D1nrsZ4FYHViCw8v%2BAGEY/U%3D\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = forked_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3eb81d1f240d926",
   "metadata": {},
   "source": [
    "# Want to know more?\n",
    "\n",
    "[Link to video walking through this notebook](https://youtu.be/hqutVJyd3TI).\n",
    "\n",
    "[https://github.com/dagworks-inc/burr](https://github.com/dagworks-inc/burr)\n",
    "<img src=\"burr_qrcode.png\" width=\"125\"/>\n",
    "\n",
    "[Time Travel blog post & video:](https://blog.dagworks.io/p/travel-back-in-time-with-burr)\n",
    "\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\"><img src=\"https://img.youtube.com/vi/98vxhIcE6NI/0.jpg\"> \n",
    "</img>\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\">\n",
    "\n",
    "More blogs @ `blog.dagworks.io` e.g. [async & streaming](https://blog.dagworks.io/p/streaming-chatbot-with-burr-fastapi)\n",
    "\n",
    "More [examples](https://github.com/DAGWorks-Inc/burr/tree/main/examples/):\n",
    "\n",
    "- e.g. [test case creation](https://burr.dagworks.io/examples/guardrails/creating_tests/)\n",
    "- e.g. [multi-agent collaboration](https://github.com/DAGWorks-Inc/burr/tree/main/examples/multi-agent-collaboration)\n",
    "\n",
    "Follow on Twitter & LinkedIn:\n",
    "\n",
    "- https://x.com/burr_framework\n",
    "- https://x.com/dagworks\n",
    "- https://x.com/stefkrawczyk\n",
    "- https://www.linkedin.com/in/skrawczyk/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54f8b9d-5782-45d6-b0e1-a91733cd6433",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
